package za.org.jcicada.comms.impl.tcp;

import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.channels.CancelledKeyException;
import java.nio.channels.ClosedChannelException;
import java.nio.channels.ClosedSelectorException;
import java.nio.channels.DatagramChannel;
import java.nio.channels.SelectableChannel;
import java.nio.channels.SelectionKey;
import java.nio.channels.Selector;
import java.nio.channels.ServerSocketChannel;
import java.nio.channels.SocketChannel;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.LinkedBlockingQueue;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;

import za.org.jcicada.comms.IAcceptor;
import za.org.jcicada.comms.impl.MessageExecutor;

/**
 * Implements NIO selector loop for handling multiple socket
 * connections from a single thread.
 * @author rcrida
 *
 */
public class TcpSelector extends Thread {

	private static final Log logger = LogFactory.getLog(TcpSelector.class);

	/** The selector object */
	private Selector selector;
	
	/** Flag used to help terminate operation of the select thread */
	private boolean running = true;
	
	/** Singleton instance of this class. */
	private static TcpSelector instance;
	
	/** Buffer for reading into. */
	private ByteBuffer readBuffer = ByteBuffer.allocate(8192);
	
	/** A list of ChangeRequest instances. */
	private BlockingQueue<ChangeRequest> changeRequests = new LinkedBlockingQueue<ChangeRequest>();

	/** Maps a SelectableChannel to a list of ByteBuffer instances. */
	private Map<SelectableChannel, List<ByteBuffer>> pendingData = new HashMap<SelectableChannel, List<ByteBuffer>>();
	
	public enum Type { REGISTER, CHANGE_OPS, ADD_OPS };
	
	public interface ICallback {
		int read(ByteBuffer dst) throws IOException;
		int write(ByteBuffer src) throws IOException;
		void close() throws IOException;
	}
	
	/** Used to change the interested operations associated with a channel. */
	private class ChangeRequest {
		private SelectableChannel channel;
		private Type type;
		private int ops;
		private Object attachment;
		
		public ChangeRequest(SelectableChannel channel, Type type, int ops) {
			this.channel = channel;
			this.type = type;
			this.ops = ops;
		}
		
		public ChangeRequest(SelectableChannel channel, Type type, int ops, Object attachment) {
			this.channel = channel;
			this.type = type;
			this.ops = ops;
			this.attachment = attachment;
		}
		
		public void apply() {
			SelectionKey key;
			switch (type) {
			case REGISTER:
				try {
					channel.register(selector, ops, attachment);
				} catch (ClosedChannelException ex) {
					logger.warn(ex);
				} catch (NullPointerException ex) {
					logger.error(ex);
				}
				break;
			case CHANGE_OPS:
				key = channel.keyFor(selector);
				try {
					key.interestOps(ops);
				} catch (CancelledKeyException ex) {
					logger.warn(ex);
				} catch (NullPointerException ex) {
					logger.error(ex);
				}
				break;
			case ADD_OPS:
				key = channel.keyFor(selector);
				try {
					key.interestOps(key.interestOps() | ops);
				} catch (CancelledKeyException ex) {
					logger.warn(ex);
				} catch (NullPointerException ex) {
					logger.error("message", ex);
				}
				break;
			}
		}
	}
	
	/** 
	 * Singleton accessor method creates instance as necessary.
	 * @throws IOException if fails to create selector
	 */
	synchronized public static TcpSelector getInstance() throws IOException {
		if (instance == null) {
			instance = new TcpSelector();
			instance.start();
		}
		return instance;
	}
	
	/**
	 * Private constructor as part of singleton pattern. Makes a daemon thread
	 * so that the JVM can exit without this thread being terminated.
	 * @throws IOException if fails to create selector
	 */
	private TcpSelector() throws IOException {
		super("Thread-" + TcpSelector.class.getSimpleName());
		setDaemon(true);
		selector = Selector.open();
	}
	
	/**
	 * Add a new acceptor to the select loop to be polled for accept events.
	 * @param acceptor the new acceptor
	 */
	public void addAcceptor(SelectableChannel channel, IAcceptor acceptor) {
		logger.info(String.format("Adding %s as acceptor", acceptor));
		changeRequests.add(new ChangeRequest(channel, Type.REGISTER, SelectionKey.OP_ACCEPT, acceptor));
		selector.wakeup();
		Thread.yield();
	}
	
	/** 
	 * Add a new tcp reader to the select loop.
	 * @param connection the new reader connection
	 */
	public void addReader(SocketChannel channel, TcpConnection connection) {
		int ops = SelectionKey.OP_READ;
		if (channel.isConnectionPending()) {
			logger.info("Connecting " + connection);
			ops |= SelectionKey.OP_CONNECT;
		}
		addReader(channel, ops, connection);
	}
	
	/**
	 * Add a new udp reader to the select loop.
	 */
	public void addReader(DatagramChannel channel, ICallback callback) {
		addReader(channel, SelectionKey.OP_READ, callback);
	}
	
	private void addReader(SelectableChannel channel, int ops, Object attachment) {
		logger.info(String.format("Adding %s as reader", attachment));
		changeRequests.add(new ChangeRequest(channel, Type.REGISTER, ops, attachment));
		selector.wakeup();
		Thread.yield();
	}
		
	/**
	 * Send data using the specified socket channel. Queues the data to be sent
	 * by within the selector thread.
	 * 
	 * @param channel
	 *            the socket channel to send on
	 * @param data
	 *            the data to be sent
	 */
	public void send(SelectableChannel channel, ByteBuffer data) {
		// And queue the data we want written
		synchronized (pendingData) {
			List<ByteBuffer> queue = pendingData.get(channel);
			if (queue == null) {
				queue = new ArrayList<ByteBuffer>();
				pendingData.put(channel, queue);
			}
			queue.add(data);
		}
		// Indicate we want the interest ops set changed
		changeRequests.add(new ChangeRequest(channel, Type.ADD_OPS, SelectionKey.OP_WRITE));

		// Finally, wake up our selecting thread so it can make the required changes
		selector.wakeup();
	}

	/**
	 * Close the selector and terminates the select thread. Calls close
	 * on selector to ensure that it terminates immediately.
	 * @throws IOException
	 */
	public void close() throws IOException {
		running = false;
		selector.close();
		try {
			join();
		} catch (InterruptedException e) {
			logger.error(e);
		}
		// destroy instance so next request makes a new one
		synchronized (TcpSelector.class) {
			TcpSelector.instance = null;
		}
        MessageExecutor.shutdownNow();
	}
	
	/**
	 * Run method contains the actual select loop. This repeats indefinitely until
	 * the running flag is cleared or an IOException occurs.
	 */
	@Override
	public void run() {
		try {
			while (running) {
				registerChanges();
				
				// wait for an event on one of the registered channels
				selector.select();
				
				try {
					// iterate over keys for which events are available
					Iterator<SelectionKey> readyIterator = selector.selectedKeys().iterator();
					while (readyIterator.hasNext()) {
						SelectionKey key = (SelectionKey) readyIterator.next();
						try {
							if (key.isAcceptable()) {
								accept(key);
							}
							if (key.isConnectable()) {
								connect(key);
							}
							if (key.isReadable()) {
								read(key);
							}
							if (key.isWritable()) {
								write(key);
							}
						} catch (IOException ex) {
							logger.error("Error with channel, cancelling key for " + key.attachment(), ex);
							cancelKey(key);
						} catch (CancelledKeyException ex) {
							// probably just been cancelled
						}
						readyIterator.remove();
					}
				} catch (CancelledKeyException ex) {
					logger.warn("Cancelled key exception", ex);
				}
			}
		} catch (IOException ex) {
			logger.fatal("Selector failure", ex);
		} catch (ClosedSelectorException ex) {
			logger.info("Selector closed");
		}
	}

	/**
	 * Iterate over registered changes and apply them to the selection keys.
	 */
	private void registerChanges() {
		// Process any pending changes
		ChangeRequest change;
		while ((change = changeRequests.poll()) != null) {
			change.apply();
		}
	}

	/**
	 * Accept a connection on the specified key.
	 * 
	 * @param key
	 * @throws IOException
	 */
	private void accept(SelectionKey key) throws IOException {
		ServerSocketChannel serverSocketChannel = (ServerSocketChannel) key.channel();
		SocketChannel socketChannel = serverSocketChannel.accept();
		socketChannel.configureBlocking(false);
		try {
			((TcpAcceptor) key.attachment()).accept(socketChannel);
		} catch (Exception ex) {
			logger.error("Problem accepting connection", ex);
		}
	}
	
	private void connect(SelectionKey key) throws IOException {
		SocketChannel socketChannel = (SocketChannel) key.channel();

		// Finish the connection. If the connection operation failed
		// this will raise an IOException.
		try {
			socketChannel.finishConnect();
		} catch (IOException e) {
			logger.error("Connection error", e);
			// Cancel the channel's registration with our selector
			cancelKey(key);
			try {
				((TcpConnection) key.attachment()).close();
			} catch (IOException e1) {
				logger.error("Problem closing connecting connection", e1);
			}
			return;
		}

		// Register an interest in reading and  writing on this channel
		key.interestOps(SelectionKey.OP_READ | SelectionKey.OP_WRITE);
	}
	
	/**
	 * The specified key indicates a channel ready to read.
	 * 
	 * @param key
	 * @throws IOException
	 */
	private void read(SelectionKey key) throws IOException {
		ICallback callback = (ICallback) key.attachment();
		int numRead;
		try {
			numRead = callback.read(readBuffer);
		} catch (IOException ex) {
			logger.info("Remote forcibly closed connection.", ex);
			cancelKey(key);
			try {
				callback.close();
			} catch (IOException ex2) {
				logger.error("Error closing listener", ex2);
			}
			return;
		}
		if (numRead == -1) {
			logger.info("EOF for " + callback);
			cancelKey(key);
			try {
				callback.close();
			} catch (IOException ex) {
				logger.error("Error closing listener", ex);
			}
			return;
		}
	}
	
	/**
	 * The key indicates a channel ready to write.
	 * 
	 * @param key
	 * @throws IOException
	 */
	private void write(SelectionKey key) throws IOException {
		SelectableChannel channel = key.channel();
		ICallback callback = (ICallback) key.attachment();

		synchronized (pendingData) {
			List<ByteBuffer> queue = pendingData.get(channel);

			if (queue != null) {
				// Write until there's not more data ...
				while (!queue.isEmpty()) {
					ByteBuffer buf = (ByteBuffer) queue.get(0);
					callback.write(buf);
					if (buf.remaining() > 0) {
						// ... or the socket's buffer fills up
						break;
					}
					queue.remove(0);
				}
	
				if (queue.isEmpty()) {
					// We wrote away all data, so we're no longer interested
					// in writing on this socket. Switch back to waiting for
					// data.
					key.interestOps(SelectionKey.OP_READ);
				}
			}
		}
	}
	
	private void cancelKey(SelectionKey key) {
		synchronized (pendingData) {
			pendingData.remove(key);
		}
		key.cancel();
	}

}
